import asyncio
import os
import dotenv

from openai import AsyncOpenAI

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    InputGuardrail,GuardrailFunctionOutput,
    function_tool,
    set_tracing_disabled,
)
from pydantic import BaseModel

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()
class HomeworkOutput(BaseModel):
    is_homework: bool
    reasoning: str

guardrail_agent = Agent(
    name="防护栏检查",
    instructions="""
你是一个作业类型检查助手。请判断用户输入是否为历史或数学相关的问题。如果是，请输出：
{"is_homework": true, "reasoning": "你的理由"}
否则输出：
{"is_homework": false, "reasoning": "你的理由"}
请只输出JSON，不要输出其他内容。
""",
    output_type=HomeworkOutput,
)
async def homework_guardrail(ctx, agent, input_data):
    print(f"input_data: {input_data}")
    result = await Runner.run(starting_agent=guardrail_agent,
                              input=input_data, context=ctx.context,
                              run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER))
    final_output = result.final_output_as(HomeworkOutput)
    return GuardrailFunctionOutput(
        output_info=final_output,
        tripwire_triggered=not final_output.is_homework,
    )
# 定义多个agent
history_tutor_agent   = Agent(
    name="历史辅导员",
    instructions="""
    你是一个历史辅导员，你的任务只是帮助用户解决历史问题。
    每一步都要解释你的推理过程并包含例子。
    如果用户的问题是历史问题，请使用历史知识来解决。
    如果用户的问题是其他问题，请回答他臣妾做不到啊.
    你在每次回答前加上 历史辅导员：
    """,
    handoff_description="""
    专门处理历史问题的 agent
    """
)
math_tutor_agent  = Agent(
    name="数学辅导员",
    handoff_description="专门处理数学问题的 agent",
    instructions="你是一个数学辅导员，你的任务是只帮助用户解决数学问题。每一步都要解释你的推理过程并包含例子。你在每次回答前加上 数学辅导员：",
)

#分类处理 Agent
triage_agent = Agent(
    name="分类 Agent",
    instructions="""
你是一个分类助手。请根据用户输入的问题内容，判断是历史问题还是数学问题。
- 如果是历史问题，只能交给history_tutor_agent 。
- 如果是数学问题，只能交给math_tutor_agent 。
- 不要自己回答，只负责分配。
- 你的分配必须准确，不能随意。
""",
    handoffs=[history_tutor_agent , math_tutor_agent ],
    input_guardrails=[
        InputGuardrail(guardrail_function=homework_guardrail),
    ],
)
async def main():
    user_input = input("请你要问的问题: ")
    result = await Runner.run(
        starting_agent = triage_agent,
        input =user_input,
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    )
    print(result.final_output)

if __name__ == "__main__":
    asyncio.run(main())